from djitellopy import Tello
from util import *


class drone_by_tello_sdk:
    def __init__(self):
        self.tl_drone = Tello()
        self.tl_drone.connect()
        self.send_rc_control = False

        self.frame_read = None

        self.moving_speed = [25, 50, 75, 100]
        self.angular_speed = [70, 80, 90, 100]
        self.ascent_speed = [30, 50, 75, 100]
        self.gear = 0

    def increase_gear(self):
        if self.gear < 3:
            self.gear += 1

    def decrease_gear(self):
        if self.gear > 0:
            self.gear -= 1

    def get_gear(self):
        return self.gear + 1

    def get_battery(self):
        return self.tl_drone.get_battery()

    @fn_timer
    def drone_move(self, a, b, c, d, speed_magnification=1):
        a_speed = -self.moving_speed[self.gear] * a * speed_magnification
        b_speed = self.moving_speed[self.gear] * b * speed_magnification
        c_speed = self.ascent_speed[self.gear] * c
        d_speed = -self.angular_speed[self.gear] * d
        a_speed = min(max(a_speed, -100), 100)
        b_speed = min(max(b_speed, -100), 100)
        c_speed = min(max(c_speed, -100), 100)
        d_speed = min(max(d_speed, -100), 100)
        if self.send_rc_control:
            self.tello.send_rc_control(self.left_right_velocity, self.for_back_velocity,
                                       self.up_down_velocity, self.yaw_velocity)

    # 起飞
    @fn_timer
    def drone_takeoff(self):
        self.tl_drone.takeoff()
        self.send_rc_control = True

    # 停止
    @fn_timer
    def drone_stop(self):
        self.tl_flight.rc(a=0, b=0, c=0, d=0)

    # 降落
    @fn_timer
    def drone_land(self):
        self.tl_drone.land()
        self.send_rc_control = False

    def drone_init_camera(self):
        self.tl_drone.streamoff()
        self.tl_drone.streamon()
        self.frame_read = self.tl_drone.get_frame_read()

    def drone_get_img(self):
        return self.frame_read.frame

    def drone_close_camera(self):
        self.tl_drone.streamoff()

    def close(self):
        self.tl_drone.streamoff()
        self.tl_drone.end()


if __name__ == '__main__':
    d = drone_by_tello_sdk()
    battery = d.get_battery()
    log('battery', battery)
    d.close()
